#include "spill/bb_spill.h"
#include "generic/logger.h"
#include "generic/random.h"
#include "generic/stats.h"
#include "basic/register.h"
#include "basic/data_dep.h"
#include "list_sched/list_sched.h"
#include "relaxed/relaxed_sched.h"
#include "enum/enumerator.h"

namespace opt_sched {

// The denominator used when calculating cost weight.
static const int COST_WGHT_BASE = 100;

BBWithSpill::BBWithSpill(MachineModel* machMdl,
                         DataDepGraph* dataDepGraph,
                         long rgnNum,
                         int16_t sigHashSize,
                         LB_ALG lbAlg,
                         LISTSCHED_HEURISTIC hurstc,
                         LISTSCHED_HEURISTIC enumHurstc,
                         bool vrfySched,
                         Pruning prune,
                         bool enblStallEnum,
                         int spillCostFactor):
  SchedRegion(machMdl, dataDepGraph, rgnNum, sigHashSize,
              lbAlg, hurstc, enumHurstc, vrfySched, prune) {
  int16_t i;

  costLwrBound_ = 0;
  enumrtr_ = NULL;
  optmlSpillCost_ = INVALID_VALUE;

  crntCycleNum_ = INVALID_VALUE;
  crntSlotNum_ = INVALID_VALUE;
  crntSpillCost_ = INVALID_VALUE;

  enblStallEnum_ = enblStallEnum;
  spillCostFactor_ = spillCostFactor;
  schedCostFactor_ = COST_WGHT_BASE;

  regTypeCnt_ = machMdl->GetRegTypeCnt();
  regFiles_ = new RegisterFile[regTypeCnt_];
  liveRegs_ = new BitVector[regTypeCnt_];
  livePhysRegs_ = new BitVector[regTypeCnt_];

  for (i = 0; i < regTypeCnt_; i++) {
    regFiles_[i].SetRegType(i);
  }
}
/****************************************************************************/

BBWithSpill::~BBWithSpill() {
  if (enumrtr_ != NULL) {
    delete enumrtr_;
  }

  delete[] regFiles_;
  delete[] liveRegs_;
  delete[] livePhysRegs_;
}
/*****************************************************************************/

ListScheduler* BBWithSpill::AllocLstSchdulr_() {
  ListScheduler* lstSchdulr = new ListScheduler(
      dataDepGraph_, machMdl_, abslutSchedUprBound_, hurstc_);
  if (lstSchdulr == NULL) Logger::Fatal("Out of memory.");
  return lstSchdulr;
}
/*****************************************************************************/

FUNC_RESULT BBWithSpill::BuildFromFile() {
  dataDepGraph_->CountDefs(regFiles_);
  dataDepGraph_->AddDefsAndUses(regFiles_);
  dataDepGraph_->AddOutputEdges();

  for (int i = 0; i < regTypeCnt_; i++) {
    liveRegs_[i].Construct(regFiles_[i].GetRegCnt());
  }

  return RES_SUCCESS;
}
/*****************************************************************************/

void BBWithSpill::SetupPhysRegs_() {
  int physRegCnt;
  for (int i = 0; i < regTypeCnt_; i++) {
    physRegCnt = regFiles_[i].FindPhysRegCnt();
    if (physRegCnt > 0)
      livePhysRegs_[i].Construct(physRegCnt);
  }
}
/*****************************************************************************/

void BBWithSpill::CmputAbslutUprBound_() {
  abslutSchedUprBound_ = dataDepGraph_->GetAbslutSchedUprBound();
  dataDepGraph_->SetAbslutSchedUprBound(abslutSchedUprBound_);
}
/*****************************************************************************/

void BBWithSpill::CmputSchedUprBound_() {
  //The maximum increase in sched length that might result in a smaller cost
  //than the known one
  int maxLngthIncrmnt = (bestCost_ - 1) / schedCostFactor_;

  assert(maxLngthIncrmnt >= 0);

  //Any schedule longer than this will have a cost that is greater than or
  //equal to that of the list schedule
  schedUprBound_ = schedLwrBound_ + maxLngthIncrmnt;

  if (abslutSchedUprBound_ < schedUprBound_) {
    schedUprBound_ = abslutSchedUprBound_;
  }
}
/*****************************************************************************/

InstCount BBWithSpill::CmputCostLwrBound() {
  // For now assume that the spill cost lower bound is 0. May be improved later
  return schedLwrBound_ * schedCostFactor_;
}
/*****************************************************************************/

void BBWithSpill::InitForSchdulng() {
  InitForCostCmputtn_();
}
/*****************************************************************************/

void BBWithSpill::InitForCostCmputtn_() {
  int i;

  crntCycleNum_ = 0;
  crntSlotNum_ = 0;
  crntSpillCost_ = INVALID_VALUE;

  for (i = 0; i < regTypeCnt_; i++) {
    regFiles_[i].ResetCrntUseCnts();
  }

  // Should set this to the live-in set
  for (i = 0; i < regTypeCnt_; i++) {
    liveRegs_[i].Reset();
    if (regFiles_[i].GetPhysRegCnt() > 0)
      livePhysRegs_[i].Reset();
  }
}
/*****************************************************************************/

InstCount BBWithSpill::CmputNormCost_(InstSchedule* sched, COST_COMP_MODE compMode,
                                        InstCount& execCost) {
  InstCount cost = CmputCost_(sched, compMode, execCost);

  cost -= costLwrBound_;
  execCost -= costLwrBound_;

  sched->SetCost(cost);
  sched->SetExecCost(execCost);
  return cost;
}
/*****************************************************************************/

InstCount BBWithSpill::CmputCost_(InstSchedule* sched, COST_COMP_MODE compMode,
                                  InstCount& execCost) {
/*  InstCount instNum;
  InstCount cycleNum;
  InstCount slotNum;
  SchedInstruction* inst;

  if (compMode == CCM_STTC) {
    InitForCostCmputtn_();

    for (instNum = sched->GetFrstInst(cycleNum, slotNum);
         instNum != INVALID_VALUE;
         instNum = sched->GetNxtInst(cycleNum, slotNum)) {
      inst = dataDepGraph_->GetInstByIndx(instNum);
      SchdulInst(inst, cycleNum, slotNum);
    }
  }*/

  assert(sched->IsComplete());
  InstCount cost = sched->GetCrntLngth() * schedCostFactor_;
  execCost = cost;
  cost += crntSpillCost_ * spillCostFactor_;
  return cost;
}
/*****************************************************************************/

void BBWithSpill::UpdateSpillInfoForSchdul_(SchedInstruction* inst) {
  int16_t regType;
  int defCnt, useCnt, regNum, physRegNum;
  Register** defs, **uses;
  Register* def, *use;
  int excessRegs, liveRegs;
  InstCount newSpillCost;

  defCnt = inst->GetDefs(defs);
  useCnt = inst->GetUses(uses);

  // Update Live regs
  for (int i = 0; i < defCnt; i++) {
    def = defs[i];
    regType = def->GetType();
    regNum = def->GetNum();
    physRegNum = def->GetPhysicalNumber(); 
    
    if (def->GetUseCnt() > 0) {
      liveRegs_[regType].SetBit(regNum, true);
      if (regFiles_[regType].GetPhysRegCnt() > 0 && physRegNum >= 0)
         livePhysRegs_[regType].SetBit(physRegNum, true);
      def->ResetCrntUseCnt();
    }
  }

  for (int i = 0; i < useCnt; i++) {
    use = uses[i];
    regType = use->GetType();
    regNum = use->GetNum();
    physRegNum = use->GetPhysicalNumber(); 

    use->AddCrntUse();

    if (use->IsLive() == false) {
      liveRegs_[regType].SetBit(regNum, false);
      if (regFiles_[regType].GetPhysRegCnt() > 0 && physRegNum >= 0)
        livePhysRegs_[regType].SetBit(physRegNum, false);
    }
  }

  newSpillCost = 0;

  for (int16_t i = 0; i < regTypeCnt_; i++) {
    liveRegs = liveRegs_[i].GetOneCnt();
    excessRegs = liveRegs - machMdl_->GetPhysRegCnt(i);

    if (excessRegs > 0) {
      newSpillCost += excessRegs;
    }
  }

  if (crntSpillCost_ == INVALID_VALUE || newSpillCost > crntSpillCost_) {
    crntSpillCost_ = newSpillCost;
  }

//Logger::Info("Scheduling inst %d: live regs = %d, crntSpillCost = %d.",
//             inst->GetNum(), liveRegs, crntSpillCost_);

}
/*****************************************************************************/

void BBWithSpill::UpdateSpillInfoForUnSchdul_(SchedInstruction* inst) {
  int16_t regType;
  int i, defCnt, useCnt, regNum, physRegNum;
  Register** defs, **uses;
  Register* def, *use;
  bool isLive;

  defCnt = inst->GetDefs(defs);
  useCnt = inst->GetUses(uses);

  // Update Live regs
  for (i = 0; i < defCnt; i++) {
    def = defs[i];
    regType = def->GetType();
    regNum = def->GetNum();
    physRegNum = def->GetPhysicalNumber();

    if (def->GetUseCnt() > 0) {
      assert(liveRegs_[regType].GetBit(regNum));
      liveRegs_[regType].SetBit(regNum, false);
      if (regFiles_[regType].GetPhysRegCnt() > 0 && physRegNum >= 0)
        livePhysRegs_[regType].SetBit(physRegNum, false);
      def->ResetCrntUseCnt();
    }
  }

  for (i = 0; i < useCnt; i++) {
    use = uses[i];
    regType = use->GetType();
    regNum = use->GetNum();
    physRegNum = use->GetPhysicalNumber();

    isLive = use->IsLive();
    use->DelCrntUse();
    assert(use->IsLive());

    if (isLive == false) {
      liveRegs_[regType].SetBit(regNum, true);
      if (regFiles_[regType].GetPhysRegCnt() > 0 && physRegNum >= 0)
        livePhysRegs_[regType].SetBit(physRegNum, true);
    }
  }

  //Logger::Info("Unscheduling inst %d: live regs = %d.",
  //             inst->GetNum(), liveRegs_[0].GetOneCnt());
}
/*****************************************************************************/

void BBWithSpill::SchdulInst(SchedInstruction* inst,
                             InstCount cycleNum,
                             InstCount slotNum) {
  crntCycleNum_ = cycleNum;
  crntSlotNum_ = slotNum;
  if (inst == NULL) return;
  assert(inst != NULL);
  UpdateSpillInfoForSchdul_(inst);
}
/*****************************************************************************/

void BBWithSpill::UnschdulInst(SchedInstruction* inst, InstCount cycleNum,
                               InstCount slotNum, EnumTreeNode* trgtNode) {
  if (slotNum == 0) {
    crntCycleNum_ = cycleNum - 1;
    crntSlotNum_ = machMdl_->GetIssueRate() - 1;
  } else {
    crntCycleNum_ = cycleNum;
    crntSlotNum_ = slotNum - 1;
  }

  if (inst == NULL) {
    return;
  }

  assert(inst != NULL);

  UpdateSpillInfoForUnSchdul_(inst);
  crntSpillCost_ = trgtNode->GetCost();

  //Logger::Info("crntSpillCost = %d.", crntSpillCost_);
}
/*****************************************************************************/

void BBWithSpill::FinishHurstc_() {
#ifdef IS_DEBUG_BBSPILL_COST
  Stats::traceCostLowerBound.Record(costLwrBound_);
  Stats::traceHeuristicCost.Record(hurstcCost_);
  Stats::traceHeuristicScheduleLength.Record(hurstcSchedLngth_);
#endif
}
/*****************************************************************************/

void BBWithSpill::FinishOptml_() {
#ifdef IS_DEBUG_BBSPILL_COST
  Stats::traceOptimalCost.Record(bestCost_);
  Stats::traceOptimalScheduleLength.Record(bestSchedLngth_);
#endif
}
/*****************************************************************************/

Enumerator* BBWithSpill::AllocEnumrtr_(Milliseconds timeout) {
  bool enblStallEnum = enblStallEnum_;

  if (!dataDepGraph_->IncludesUnpipelined()) {
    enblStallEnum = false;
  }

  enumrtr_ = new LengthCostEnumerator(dataDepGraph_, machMdl_,
                                      schedUprBound_, sigHashSize_,
                                      enumHurstc_, prune_, enblStallEnum,
                                      timeout, 0, NULL);
  if (enumrtr_ == NULL) Logger::Fatal("Out of memory.");

  return enumrtr_;
}
/*****************************************************************************/

FUNC_RESULT BBWithSpill::Enumerate_(Milliseconds deadline) {
  InstCount   trgtLngth;
  FUNC_RESULT   rslt = RES_SUCCESS;
  int           iterCnt = 0;
  int           costLwrBound = 0;

  for (trgtLngth = schedLwrBound_; trgtLngth <= schedUprBound_; trgtLngth++) {
    InitForSchdulng();
    rslt = enumrtr_->FindFeasibleSchedule(
        enumCrntSched_, trgtLngth, this, costLwrBound, deadline);
    HandlEnumrtrRslt_(rslt, trgtLngth);

    if (bestCost_ == 0 || rslt == RES_TIMEOUT || rslt == RES_ERROR) break;

    enumrtr_->Reset();
    CmputSchedUprBound_();
    iterCnt++;
    costLwrBound += 1;
  }

#ifdef IS_DEBUG_ITERS
  Stats::iterations.Record(iterCnt);
  Stats::enumerations.Record(enumrtr_->GetSearchCnt());
  Stats::lengths.Record(iterCnt);
#endif

  //Failure to find a feasible sched. in the last iteration is still
  //considered an overall success
  if (rslt == RES_SUCCESS || rslt == RES_FAIL) {
    rslt = RES_SUCCESS;
  }

  return rslt;
}
/*****************************************************************************/

InstCount BBWithSpill::UpdtOptmlSched(InstSchedule* crntSched,
                                      LengthCostEnumerator*) {
  InstCount crntCost;
  InstCount crntExecCost;

  crntCost = CmputNormCost_(crntSched, CCM_DYNMC, crntExecCost);

#ifdef IS_DEBUG_SOLN_DETAILS_2
  Logger::Info("Found a feasible sched. of length %d, cost %d",
               crntSched->GetCrntLngth(), crntCost);
  crntSched->Print(Logger::GetLogStream(), " ");
#endif

  if (crntCost < bestCost_) {
    bestCost_ = crntCost;
    optmlSpillCost_ = crntSpillCost_;
    bestSchedLngth_ = crntSched->GetCrntLngth();
    bestMainSchedLngth_ = crntSched->GetCrntLngth();
    enumBestSched_->Copy(crntSched);
  }

  return bestCost_;
}
/*****************************************************************************/

void BBWithSpill::SetupForSchdulng_() {
  SetupPhysRegs_();
}
/*****************************************************************************/

bool BBWithSpill::ChkCostFsblty(InstCount trgtLngth,
                                EnumTreeNode*,
                                SchedInstruction*,
                                LinkedList<SchedInstruction>*,
                                bool,
                                InstCount& cost,
                                InstCount& dynmcCostLwrBound,
                                LengthCostEnumerator*) {
  bool fsbl = true;

  cost = crntSpillCost_;
  dynmcCostLwrBound = crntSpillCost_ * spillCostFactor_ + trgtLngth * schedCostFactor_;
  dynmcCostLwrBound -= costLwrBound_;

  assert(cost >= 0);
  assert(dynmcCostLwrBound >= 0);

  fsbl = dynmcCostLwrBound < bestCost_;

  return fsbl;
}
/*****************************************************************************/

void BBWithSpill::SetSttcLwrBounds(EnumTreeNode*) {
  // Nothing.
}

/*****************************************************************************/

bool BBWithSpill::ChkInstLglty(SchedInstruction *inst) {
  int16_t regType;
  int defCnt, physRegNum;
  Register** defs;
  Register *def, *liveDef;

  defCnt = inst->GetDefs(defs);

  // Update Live regs
  for (int i = 0; i < defCnt; i++) {
    def = defs[i];
    regType = def->GetType();
    physRegNum = def->GetPhysicalNumber(); 
    
    // If this is a physical register definition and another
    // definition of the same physical register is live, then
    // scheduling this instruction is illegal unless this
    // instruction is the last use of that physical reg definition.
    if (regFiles_[regType].GetPhysRegCnt() > 0 && 
        physRegNum >= 0 &&
        livePhysRegs_[regType].GetBit(physRegNum) == true) {

        liveDef = regFiles_[regType].FindLiveReg(physRegNum);
        assert (liveDef != NULL);

        // If this instruction is the last use of the current live def
        if (liveDef->GetCrntUseCnt()+1 == liveDef->GetUseCnt() && inst->FindUse(liveDef)==true)
          return true;
        else
          return false;
    } // end if     
  } // end for
  return true;
}

} // end namespace opt_sched
